"""
人脸识别模型训练脚本

该脚本使用InceptionResnetV1模型和Triplet Loss进行人脸识别模型训练。
支持数据增强、模型微调、训练/验证集划分等功能。
"""

import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset, random_split
from torchvision import transforms
from facenet_pytorch import InceptionResnetV1, MTCNN
import cv2
import numpy as np
from PIL import Image
import argparse
# sklearn.model_selection模块用于数据集划分
from sklearn.model_selection import train_test_split
import random

class TripletLoss(nn.Module):
    """
    Triplet Loss函数实现
    
    Triplet Loss是一种用于学习特征嵌入的损失函数，通过使同类样本特征更接近，
    异类样本特征更远离来提升特征的判别性。
    """
    
    def __init__(self, margin=0.2):
        """
        初始化Triplet Loss
        
        参数:
            margin (float): 间隔参数，控制正负样本间的最小距离差
        """
        super(TripletLoss, self).__init__()
        self.margin = margin

    def forward(self, anchor, positive, negative):
        """
        前向传播计算Triplet Loss
        
        参数:
            anchor (Tensor): anchor样本特征向量
            positive (Tensor): positive样本特征向量（与anchor同类）
            negative (Tensor): negative样本特征向量（与anchor异类）
            
        返回:
            Tensor: 平均损失值
        """
        # 计算anchor与positive之间的欧氏距离平方
        distance_positive = torch.pow(torch.nn.functional.pairwise_distance(anchor, positive, keepdim=True), 2)
        # 计算anchor与negative之间的欧氏距离平方
        distance_negative = torch.pow(torch.nn.functional.pairwise_distance(anchor, negative, keepdim=True), 2)
        # 计算损失：max(0, d(a,p) - d(a,n) + margin)
        losses = torch.relu(distance_positive - distance_negative + self.margin)
        # 返回平均损失
        return torch.mean(losses)

class FaceDataset(Dataset):
    """
    人脸数据集类
    
    用于加载人脸图像数据，支持按人名组织的文件夹结构。
    每个子文件夹代表一个人，文件夹名即为人名标签。
    """
    
    def __init__(self, root_dir, transform=None):
        """
        初始化人脸数据集
        
        参数:
            root_dir (str): 数据集根目录路径
            transform (callable, optional): 图像变换操作
        """
        self.root_dir = root_dir
        self.transform = transform
        self.images = []  # 图像路径列表
        self.labels = []  # 标签列表
        self.label_to_name = {}  # 标签到人名的映射
        self.name_to_label = {}  # 人名到标签的映射
        
        # 加载图像和标签
        self._load_dataset()
        
    def _load_dataset(self):
        """
        从文件系统加载数据集
        
        遍历根目录下的所有子目录，每个子目录代表一个人，
        子目录名作为人名标签，目录中的图像文件作为训练样本。
        """
        label_counter = 0
        for person_name in os.listdir(self.root_dir):
            person_dir = os.path.join(self.root_dir, person_name)
            if not os.path.isdir(person_dir):
                continue
                
            # 为新人名分配标签
            if person_name not in self.name_to_label:
                self.name_to_label[person_name] = label_counter
                self.label_to_name[label_counter] = person_name
                label_counter += 1
                
            # 遍历人名目录中的所有图像文件
            for image_file in os.listdir(person_dir):
                if image_file.lower().endswith(('.jpg', '.jpeg', '.png')):
                    image_path = os.path.join(person_dir, image_file)
                    self.images.append(image_path)
                    self.labels.append(self.name_to_label[person_name])
    
    def __len__(self):
        """
        返回数据集大小
        
        返回:
            int: 数据集中图像的总数
        """
        return len(self.images)
    
    def __getitem__(self, idx):
        """
        获取指定索引的数据样本
        
        参数:
            idx (int): 样本索引
            
        返回:
            tuple: (图像, 标签) 元组
        """
        image_path = self.images[idx]
        label = self.labels[idx]
        
        # 读取图像
        image = cv2.imread(image_path)
        # 转换颜色空间从BGR到RGB
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        # 转换为PIL图像格式
        image = Image.fromarray(image)
        
        # 应用图像变换
        if self.transform:
            image = self.transform(image)
            
        return image, label

class FaceRecognitionTrainer:
    """
    人脸识别模型训练器
    
    负责模型初始化、数据加载、训练循环和模型保存等完整训练流程。
    """
    
    def __init__(self, data_dir, model_save_path="best_model.pth", batch_size=8):
        """
        初始化训练器
        
        参数:
            data_dir (str): 人脸数据集目录路径
            model_save_path (str): 最佳模型保存路径
            batch_size (int): 批处理大小
        """
        self.data_dir = data_dir
        self.model_save_path = model_save_path
        self.batch_size = batch_size
        # 检查是否有可用的GPU，如果有则使用GPU，否则使用CPU
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
        # 初始化预训练的InceptionResnetV1模型
        self.model = InceptionResnetV1(
            pretrained='vggface2',  # 使用在VGGFace2数据集上预训练的权重
            classify=False,         # 使用特征提取模式，而非分类模式
            num_classes=None        # 不指定分类数量
        ).to(self.device)
        
        # 冻结底层网络参数
        self._freeze_base_layers()
        
        # 定义训练时的数据增强操作
        self.train_transform = transforms.Compose([
            transforms.RandomHorizontalFlip(p=0.5),     # 50%概率随机水平翻转
            transforms.RandomRotation(degrees=10),      # 随机旋转±10度
            transforms.ToTensor(),                      # 转换为Tensor
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # 归一化
        ])
        
        # 定义测试时的数据变换操作（无增强）
        self.test_transform = transforms.Compose([
            transforms.ToTensor(),                      # 转换为Tensor
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])  # 归一化
        ])
        
        # 加载数据集
        self._load_datasets()
        
        # 定义损失函数和优化器
        self.criterion = TripletLoss(margin=0.2)
        # 只优化未冻结层的参数
        self.optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.model.parameters()), lr=0.001)
        
        # 学习率调度器，每10个epoch将学习率乘以0.1
        self.scheduler = optim.lr_scheduler.StepLR(self.optimizer, step_size=10, gamma=0.1)
        
        # 记录最佳验证损失
        self.best_loss = float('inf')
        
    def _freeze_base_layers(self):
        """
        冻结底层网络参数
        
        冻结大部分网络层，只保留最后几个层可训练，以适应特定任务。
        """
        # 冻结大部分层，只保留最后的几个层可训练
        for name, child in self.model.named_children():
            # 保持last_bn和last_conv层可训练
            if name in ['last_bn', 'last_conv']:
                print(f"保持 {name} 层可训练")
                for param in child.parameters():
                    param.requires_grad = True
            else:
                print(f"冻结 {name} 层")
                for param in child.parameters():
                    param.requires_grad = False
    
    def _load_datasets(self):
        """
        加载训练集和测试集
        
        将数据集按8:2的比例划分为训练集和测试集。
        """
        # 创建完整数据集
        full_dataset = FaceDataset(self.data_dir, transform=self.train_transform)
        
        # 计算训练集和测试集大小
        train_size = int(0.8 * len(full_dataset))
        test_size = len(full_dataset) - train_size
        # 划分训练集和测试集
        self.train_dataset, self.test_dataset = random_split(
            full_dataset, 
            [train_size, test_size],
            generator=torch.Generator().manual_seed(42)  # 固定随机种子以确保可重复性
        )
        
        # 为测试集应用不同的变换（无数据增强）
        self.test_dataset.dataset.transform = self.test_transform
        
        # 创建训练集数据加载器
        self.train_loader = DataLoader(
            self.train_dataset, 
            batch_size=self.batch_size, 
            shuffle=True,               # 打乱数据
            num_workers=0               # 在Windows上可能需要设置为0
        )
        
        # 创建测试集数据加载器
        self.test_loader = DataLoader(
            self.test_dataset, 
            batch_size=self.batch_size, 
            shuffle=False,              # 不打乱数据
            num_workers=0
        )
        
        print(f"训练集大小: {len(self.train_dataset)}")
        print(f"测试集大小: {len(self.test_dataset)}")
        print(f"类别数量: {len(full_dataset.name_to_label)}")
    
    def _get_triplet(self, anchor_label, batch_images, batch_labels):
        """
        从批次中获取正样本和负样本
        
        构建三元组(anchor, positive, negative)用于Triplet Loss计算。
        
        参数:
            anchor_label (int): anchor样本的标签
            batch_images (Tensor): 批次中的所有图像
            batch_labels (Tensor): 批次中所有图像的标签
            
        返回:
            tuple: (positive_image, negative_image) 或 (None, None) 如果无法构建三元组
        """
        # 找到正样本索引（相同标签）
        positive_indices = torch.where(batch_labels == anchor_label)[0]
        # 找到负样本索引（不同标签）
        negative_indices = torch.where(batch_labels != anchor_label)[0]
        
        # 如果正样本少于2个或没有负样本，则无法构建三元组
        if len(positive_indices) < 2 or len(negative_indices) == 0:
            return None, None
        
        # 随机选择正样本和负样本
        positive_idx = random.choice(positive_indices)
        negative_idx = random.choice(negative_indices)
        
        return batch_images[positive_idx], batch_images[negative_idx]
    
    def train_epoch(self):
        """
        训练一个epoch
        
        执行一个完整的训练周期，包括前向传播、损失计算和反向传播。
        
        返回:
            float: 平均训练损失
        """
        self.model.train()  # 设置模型为训练模式
        total_loss = 0.0    # 累计损失
        num_batches = 0     # 批次数
        
        # 遍历训练集数据加载器
        for batch_idx, (images, labels) in enumerate(self.train_loader):
            # 将数据移动到指定设备（GPU或CPU）
            images, labels = images.to(self.device), labels.to(self.device)
            
            # 清零梯度
            self.optimizer.zero_grad()
            
            # 存储anchor, positive, negative样本
            anchors, positives, negatives = [], [], []
            
            # 为批次中的每个样本构建三元组
            for i in range(len(images)):
                anchor_img = images[i].unsqueeze(0)  # 增加批次维度
                anchor_label = labels[i]
                
                # 获取正样本和负样本
                positive_img, negative_img = self._get_triplet(anchor_label, images, labels)
                
                # 如果成功获取三元组，则添加到列表中
                if positive_img is not None and negative_img is not None:
                    anchors.append(anchor_img)
                    positives.append(positive_img.unsqueeze(0))
                    negatives.append(negative_img.unsqueeze(0))
            
            # 如果有有效的三元组
            if len(anchors) > 0:
                # 合并所有三元组
                anchors = torch.cat(anchors, dim=0)
                positives = torch.cat(positives, dim=0)
                negatives = torch.cat(negatives, dim=0)
                
                # 前向传播计算特征嵌入
                anchor_embeddings = self.model(anchors)
                positive_embeddings = self.model(positives)
                negative_embeddings = self.model(negatives)
                
                # 计算Triplet Loss
                loss = self.criterion(anchor_embeddings, positive_embeddings, negative_embeddings)
                
                # 反向传播
                loss.backward()
                self.optimizer.step()
                
                # 累计损失和批次数
                total_loss += loss.item()
                num_batches += 1
                
                # 每10个批次打印一次损失
                if batch_idx % 10 == 0:
                    print(f'训练批次 [{batch_idx}/{len(self.train_loader)}] 损失: {loss.item():.6f}')
        
        # 计算平均损失
        avg_loss = total_loss / num_batches if num_batches > 0 else 0
        return avg_loss
    
    def validate(self):
        """
        验证模型
        
        在测试集上评估模型性能。
        
        返回:
            float: 平均验证损失
        """
        self.model.eval()   # 设置模型为评估模式
        total_loss = 0.0    # 累计损失
        num_batches = 0     # 批次数
        
        # 不计算梯度，节省内存并加快计算
        with torch.no_grad():
            # 遍历测试集数据加载器
            for batch_idx, (images, labels) in enumerate(self.test_loader):
                # 将数据移动到指定设备（GPU或CPU）
                images, labels = images.to(self.device), labels.to(self.device)
                
                # 存储anchor, positive, negative样本
                anchors, positives, negatives = [], [], []
                
                # 为批次中的每个样本构建三元组
                for i in range(len(images)):
                    anchor_img = images[i].unsqueeze(0)  # 增加批次维度
                    anchor_label = labels[i]
                    
                    # 获取正样本和负样本
                    positive_img, negative_img = self._get_triplet(anchor_label, images, labels)
                    
                    # 如果成功获取三元组，则添加到列表中
                    if positive_img is not None and negative_img is not None:
                        anchors.append(anchor_img)
                        positives.append(positive_img.unsqueeze(0))
                        negatives.append(negative_img.unsqueeze(0))
                
                # 如果有有效的三元组
                if len(anchors) > 0:
                    # 合并所有三元组
                    anchors = torch.cat(anchors, dim=0)
                    positives = torch.cat(positives, dim=0)
                    negatives = torch.cat(negatives, dim=0)
                    
                    # 前向传播计算特征嵌入
                    anchor_embeddings = self.model(anchors)
                    positive_embeddings = self.model(positives)
                    negative_embeddings = self.model(negatives)
                    
                    # 计算Triplet Loss
                    loss = self.criterion(anchor_embeddings, positive_embeddings, negative_embeddings)
                    
                    # 累计损失和批次数
                    total_loss += loss.item()
                    num_batches += 1
        
        # 计算平均损失
        avg_loss = total_loss / num_batches if num_batches > 0 else 0
        return avg_loss
    
    def train(self, epochs=20):
        """
        训练模型
        
        执行完整的训练流程，包括多个epoch的训练和验证。
        
        参数:
            epochs (int): 训练轮数
        """
        print(f"开始训练模型，共 {epochs} 个epoch")
        print(f"使用设备: {self.device}")
        
        # 执行指定轮数的训练
        for epoch in range(1, epochs + 1):
            print(f"\nEpoch {epoch}/{epochs}")
            print('-' * 20)
            
            # 训练一个epoch
            train_loss = self.train_epoch()
            print(f'训练损失: {train_loss:.6f}')
            
            # 验证模型
            val_loss = self.validate()
            print(f'验证损失: {val_loss:.6f}')
            
            # 更新学习率
            self.scheduler.step()
            
            # 保存最佳模型（验证损失最低的模型）
            if val_loss < self.best_loss:
                self.best_loss = val_loss
                torch.save(self.model.state_dict(), self.model_save_path)
                print(f'保存最佳模型，验证损失: {val_loss:.6f}')
        
        print(f"\n训练完成! 最佳验证损失: {self.best_loss:.6f}")

def list_directories(path):
    """
    列出指定路径下的所有目录
    
    参数:
        path (str): 要列出目录的路径
        
    返回:
        list: 目录名称列表
    """
    if not os.path.exists(path):
        return []
    
    directories = []
    for item in os.listdir(path):
        item_path = os.path.join(path, item)
        if os.path.isdir(item_path):
            directories.append(item)
    
    return directories

def select_data_directory():
    """
    交互式选择人脸数据集目录
    
    显示face_data目录下的所有子目录，并让用户选择一个作为数据集目录。
    
    返回:
        str: 用户选择的数据集目录路径，如果未选择则返回None
    """
    face_data_path = "face_data"
    
    # 检查face_data目录是否存在
    if not os.path.exists(face_data_path):
        print(f"错误: 找不到 '{face_data_path}' 目录")
        return None
    
    # 获取所有子目录
    directories = list_directories(face_data_path)
    
    # 如果没有子目录
    if not directories:
        print(f"'{face_data_path}' 目录下没有找到任何子目录")
        return None
    
    # 显示所有可选项
    print(f"\n在 '{face_data_path}' 目录下找到以下数据集:")
    for i, directory in enumerate(directories, 1):
        print(f"{i}. {directory}")
    
    # 获取用户选择
    while True:
        try:
            choice = input(f"\n请选择数据集 (1-{len(directories)}) 或输入 'q' 退出: ").strip()
            
            if choice.lower() == 'q':
                return None
            
            choice_index = int(choice) - 1
            if 0 <= choice_index < len(directories):
                selected_dir = os.path.join(face_data_path, directories[choice_index])
                print(f"已选择数据集目录: {selected_dir}")
                return selected_dir
            else:
                print(f"无效选择，请输入 1 到 {len(directories)} 之间的数字")
        except ValueError:
            print("请输入有效的数字或 'q' 退出")

def main():
    """
    主函数
    
    通过交互式方式获取数据集目录并启动训练流程。
    """
    print("人脸识别模型训练程序")
    print("=" * 30)
    
    # 交互式选择数据集目录
    data_dir = select_data_directory()
    
    # 如果用户取消选择，则退出程序
    if data_dir is None:
        print("程序已退出")
        return
    
    # 检查选择的目录是否存在数据
    if not os.listdir(data_dir):
        print(f"错误: 选择的目录 '{data_dir}' 为空")
        return
    
    # 设置模型保存路径
    model_path = "best_model.pth"
    
    # 设置批处理大小
    batch_size = 8
    
    # 设置训练轮数
    epochs = 20
    
    print(f"\n训练配置:")
    print(f"  数据集目录: {data_dir}")
    print(f"  模型保存路径: {model_path}")
    print(f"  批处理大小: {batch_size}")
    print(f"  训练轮数: {epochs}")
    
    # 确认开始训练
    confirm = input("\n是否开始训练? (y/n): ").strip().lower()
    if confirm != 'y' and confirm != 'yes':
        print("训练已取消")
        return
    
    # 创建训练器实例
    trainer = FaceRecognitionTrainer(
        data_dir=data_dir,
        model_save_path=model_path,
        batch_size=batch_size
    )
    
    # 开始训练
    trainer.train(epochs=epochs)

# 程序入口点
if __name__ == "__main__":
    main()